# -*- coding: utf-8 -*-
import json
import time
from pathlib import Path
from typing import List, Dict

import chromadb
from llama_index.core import VectorStoreIndex, StorageContext, Settings
from llama_index.core.schema import TextNode
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import PromptTemplate

"""
使用chatgpt优化过代码，添加注释
"""

# ================== QA模板 ==================
QA_TEMPLATE = (
    "<|im_start|>system\n"
    "你是一个专业的法律助手，请严格根据以下法律条文回答问题：\n"
    "相关法律条文：\n{context_str}\n<|im_end|>\n"
    "<|im_start|>user\n{query_str}<|im_end|>\n"
    "<|im_start|>assistant\n"
)

# 设置问答模板
response_template = PromptTemplate(QA_TEMPLATE)


# ================== 配置区 ==================
class Config:
    # 设置嵌入和语言模型的路径，注释中提供了示例路径
    EMBED_MODEL_PATH = r"/智客ai/llm/BAAI/bge-small-zh-v1.5"
    LLM_MODEL_PATH = r"/智客ai/llm/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"

    # 数据目录配置
    DATA_DIR = "./data"
    VECTOR_DB_DIR = "./chroma_db"
    PERSIST_DIR = "./storage"

    # 向量数据库集合的名称
    COLLECTION_NAME = "chinese_labor_laws"

    # 设置返回的最大数量（即相关法律条文的数量）
    TOP_K = 3


# ================== 初始化模型 ==================
def init_models():
    """初始化嵌入模型（Embedding）和大语言模型（LLM）并验证其有效性"""

    # 初始化Embedding模型
    embed_model = HuggingFaceEmbedding(
        model_name=Config.EMBED_MODEL_PATH,  # 使用的嵌入模型路径
    )

    # 初始化大语言模型（LLM）
    llm = HuggingFaceLLM(
        model_name=Config.LLM_MODEL_PATH,  # 使用的LLM模型路径
        tokenizer_name=Config.LLM_MODEL_PATH,  # 使用的tokenizer路径   分词器
        model_kwargs={
            "trust_remote_code": True,  # 允许远程代码信任（用于加载模型代码）
        },
        tokenizer_kwargs={"trust_remote_code": True},  # 同样信任tokenizer代码
        generate_kwargs={"temperature": 0.3}  # 生成文本时的温度设置，控制多样性
    )

    # 将模型设置到全局设置中，方便后续访问
    Settings.embed_model = embed_model
    Settings.llm = llm

    # 验证嵌入模型，检查生成的嵌入向量维度
    test_embedding = embed_model.get_text_embedding("测试文本")
    print(f"Embedding维度验证：{len(test_embedding)}")

    return embed_model, llm


# ================== 数据处理 ==================
def load_and_validate_json_files(data_dir: str) -> List[Dict]:
    """加载并验证JSON格式的法律文件"""

    # 获取目录中所有的JSON文件
    json_files = list(Path(data_dir).glob("*.json"))
    print(f"json_files为{json_files}")
    assert json_files, f"未找到JSON文件于 {data_dir}"

    all_data = []
    for json_file in json_files:
        with open(json_file, 'r', encoding='utf-8') as f:
            try:
                # 解析JSON文件内容
                data = json.load(f)

                # 验证数据格式，确保根元素为列表，且列表项为字典
                if not isinstance(data, list):
                    raise ValueError(f"文件 {json_file.name} 根元素应为列表")
                for item in data:
                    if not isinstance(item, dict):
                        raise ValueError(f"文件 {json_file.name} 包含非字典元素")
                    for k, v in item.items():
                        if not isinstance(v, str):
                            raise ValueError(f"文件 {json_file.name} 中键 '{k}' 的值不是字符串")

                # 扩展所有数据项
                all_data.extend({
                                    "content": item,
                                    "metadata": {"source": json_file.name}
                                } for item in data)
            except Exception as e:
                # 如果读取或解析失败，抛出异常
                raise RuntimeError(f"加载文件 {json_file} 失败: {str(e)}")

    print(f"成功加载 {len(all_data)} 个法律文件条目")
    return all_data


def create_nodes(raw_data: List[Dict]) -> List[TextNode]:
    """根据加载的数据生成文本节点（TextNode）并保证ID的稳定性"""
    nodes = []
    for entry in raw_data:
        law_dict = entry["content"]  # 法律条文内容
        source_file = entry["metadata"]["source"]  # 来源文件名

        for full_title, content in law_dict.items():
            # 生成一个稳定的ID，避免重复 例如：'data1.json::中华人民共和国劳动合同法 第一条'
            node_id = f"{source_file}::{full_title}"

            parts = full_title.split(" ", 1)
            law_name = parts[0] if len(parts) > 0 else "未知法律"
            article = parts[1] if len(parts) > 1 else "未知条款"
            # 设置节点
            node = TextNode(
                text=content,
                id_=node_id,  # 显式设置稳定ID
                metadata={
                    "law_name": law_name,
                    "article": article,
                    "full_title": full_title,
                    "source_file": source_file,
                    "content_type": "legal_article"  # 表明是法律条文内容
                }
            )
            nodes.append(node)

    print(f"生成 {len(nodes)} 个文本节点（ID示例：{nodes[0].id_}）")
    return nodes


# ================== 向量存储 ==================
def init_vector_store(nodes: List[TextNode]) -> VectorStoreIndex:
    """初始化并设置向量存储（Chroma）"""

    chroma_client = chromadb.PersistentClient(path=Config.VECTOR_DB_DIR)  # 使用PersistentClient保持数据持久化
    chroma_collection = chroma_client.get_or_create_collection(
        name=Config.COLLECTION_NAME,  # 集合名称
        metadata={"hnsw:space": "cosine"}  # 设置索引空间为cosine相似度
    )

    # 创建存储上下文
    storage_context = StorageContext.from_defaults(
        vector_store=ChromaVectorStore(chroma_collection=chroma_collection)
    )

    # 判断是否需要创建新索引
    if chroma_collection.count() == 0 and nodes is not None:
        print(f"创建新索引（{len(nodes)}个节点）...")

        # 将所有文本节点添加到存储上下文中
        storage_context.docstore.add_documents(nodes)

        # 创建VectorStoreIndex
        index = VectorStoreIndex(
            nodes,
            storage_context=storage_context,
            show_progress=True  # 显示进度
        )

        # 持久化存储
        storage_context.persist(persist_dir=Config.PERSIST_DIR)
        index.storage_context.persist(persist_dir=Config.PERSIST_DIR)  # 双重持久化保障
    else:
        print("加载已有索引...")
        # 加载已存在的索引
        storage_context = StorageContext.from_defaults(
            persist_dir=Config.PERSIST_DIR,
            vector_store=ChromaVectorStore(chroma_collection=chroma_collection)
        )
        # 创建已有索引的查询引擎
        index = VectorStoreIndex.from_vector_store(
            storage_context.vector_store,
            storage_context=storage_context,
            embed_model=Settings.embed_model
        )

    # 存储验证，输出文档数目和示例ID
    print("\n存储验证结果：")
    doc_count = len(storage_context.docstore.docs)
    print(f"DocStore记录数：{doc_count}")

    if doc_count > 0:
        sample_key = next(iter(storage_context.docstore.docs.keys()))
        print(f"示例节点ID：{sample_key}")
    else:
        print("警告：文档存储为空，请检查节点添加逻辑！")

    return index


# ================== 主程序 ==================
def main():
    embed_model, llm = init_models()  # 初始化模型

    # 仅当需要更新数据时执行
    if not Path(Config.VECTOR_DB_DIR).exists():
        print("\n初始化数据...")
        # 处理数据，组装成  [{"content":{},"metadata":{"source":"data1.json"}},{"content":{},"metadata":{"source":"data1.json"}}]
        raw_data = load_and_validate_json_files(Config.DATA_DIR)  # 加载并验证数据
        nodes = create_nodes(raw_data)  # 创建文本节点
    else:
        nodes = None  # 已有数据时不加载

    print("\n初始化向量存储...")
    start_time = time.time()
    # 初始化向量存储
    index = init_vector_store(nodes)
    print(f"索引加载耗时：{time.time() - start_time:.2f}s")

    # 创建查询引擎
    query_engine = index.as_query_engine(
        similarity_top_k=Config.TOP_K,  # 设置返回最相关的前TOP_K条结果
        text_qa_template=response_template,  # 设置问答模板
        verbose=True  # 显示详细查询过程
    )

    # 示例查询，用户输入查询问题
    while True:
        question = input("\n请输入劳动法相关问题（输入q退出）: ")
        if question.lower() == 'q':
            break

        # 执行查询并获取结果
        response = query_engine.query(question)

        # 这里的响应就是response.response
        # print(f"response里面的东西有:{response}")
        """
        示例：response.source_nodes的输出为：
        [NodeWithScore(node=TextNode(id_='data1.json: :中华人民共和国劳动法 第一百条', embedding=None, metadata={'law_name': '中华人民共和国劳动法', 'article': '第一百条', 'full_title': '中华人民共和国劳动法 第一百条', 'source_file': 'data1.json', 'content_type': 'legal_article'
            }, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, metadata_template='{key
            }: {value
            }', metadata_separator='\n', text='用人单位无故不缴纳社会保险费的，由劳动行政部门责令其限期缴纳，逾期不缴的，可以加收滞纳金。', mimetype='text/plain', start_char_idx=None, end_char_idx=None, metadata_seperator='\n', text_template='{metadata_str
            }\n\n{content
            }'), score=0.7965761495434993), NodeWithScore(node=TextNode(id_='data1.json: :中华人民共和国劳动法 第七十二条', embedding=None, metadata={'law_name': '中华人民共和国劳动法', 'article': '第七十二条', 'full_title': '中华人民共和国劳动法 第七十二条', 'source_file': 'data1.json', 'content_type': 'legal_article'
            }, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, metadata_template='{key
            }: {value
            }', metadata_separator='\n', text='社会保险基金按照保险类型确定资金来源，逐步实行社会统筹。用人单位和劳动者必须依法参加社会保险，缴纳社会保险费。', mimetype='text/plain', start_char_idx=None, end_char_idx=None, metadata_seperator='\n', text_template='{metadata_str
            }\n\n{content
            }'), score=0.723787623398956), NodeWithScore(node=TextNode(id_='data1.json: :中华人民共和国劳动法 第七十三条', embedding=None, metadata={'law_name': '中华人民共和国劳动法', 'article': '第七十三条', 'full_title': '中华人民共和国劳动法 第七十三条', 'source_file': 'data1.json', 'content_type': 'legal_article'
            }, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, metadata_template='{key
            }: {value
            }', metadata_separator='\n', text='劳动者在下列情形下，依法享受社会保险待遇：\n（一）退休；\n（二）患病、负伤；\n（三）因工伤残或者患职业病；\n（四）失业；\n（五）生育。\n劳动者死亡后，其遗属依法享受遗属津贴。\n劳动者享受社会保险待遇的条件和标准由法律、法规规定。\n劳动者享受的社会保险金必须按时足额支付。', mimetype='text/plain', start_char_idx=None, end_char_idx=None, metadata_seperator='\n', text_template='{metadata_str
            }\n\n{content
            }'), score=0.6831792454535246)
        ]
        """
        print(f"response.source_nodes为:{response.source_nodes}")

        """
        示例：response.metadata的输出为：
            {'data1.json: :中华人民共和国劳动法 第一百条': {'law_name': '中华人民共和国劳动法', 'article': '第一百条', 'full_title': '中华人民共和国劳动法 第一百条', 'source_file': 'data1.json', 'content_type': 'legal_article'
                }, 'data1.json: :中华人民共和国劳动法 第七十二条': {'law_name': '中华人民共和国劳动法', 'article': '第七十二条', 'full_title': '中华人民共和国劳动法 第七十二条', 'source_file': 'data1.json', 'content_type': 'legal_article'
                }, 'data1.json: :中华人民共和国劳动法 第七十三条': {'law_name': '中华人民共和国劳动法', 'article': '第七十三条', 'full_title': '中华人民共和国劳动法 第七十三条', 'source_file': 'data1.json', 'content_type': 'legal_article'
                }
            }
        """
        print(f"response.metadata为:{response.metadata}")

        # 显示结果
        print(f"\n智能助手回答：\n{response.response}")
        print("\n支持依据：")
        for idx, node in enumerate(response.source_nodes, 1):
            meta = node.metadata  # 获取节点元数据
            print(f"\n[{idx}] {meta['full_title']}")
            print(f"  来源文件：{meta['source_file']}")
            print(f"  法律名称：{meta['law_name']}")
            print(f"  条款内容：{node.text[:100]}...")  # 显示部分条款内容
            print(f"  相关度得分：{node.score:.4f}")


if __name__ == "__main__":
    main()  # 运行主程序
